package com.zhou.base.components.feign.cloud.loadbalancer;

import com.alibaba.cloud.commons.lang.StringUtils;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultRequest;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.loadbalancer.config.LoadBalancerZoneConfig;
import org.springframework.cloud.loadbalancer.core.DelegatingServiceInstanceListSupplier;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import reactor.core.publisher.Flux;

import javax.servlet.http.HttpServletRequest;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

/**
 *
 */
public class MyZonePreferenceServiceInstanceListSupplier extends DelegatingServiceInstanceListSupplier {
    private final String ZONE = "zone";
    private final LoadBalancerZoneConfig zoneConfig;

    public MyZonePreferenceServiceInstanceListSupplier(ServiceInstanceListSupplier delegate,
                                                       LoadBalancerZoneConfig zoneConfig) {
        super(delegate);
        this.zoneConfig = zoneConfig;
    }

    @Override
    public Flux<List<ServiceInstance>> get() {
        return getDelegate().get();
    }

    private List<ServiceInstance> filteredByZone(String zone, List<ServiceInstance> serviceInstances) {
        if (zone == null) {
            zone = zoneConfig.getZone();
        }
        if (zone != null) {
            List<ServiceInstance> filteredInstances = new ArrayList<>();
            for (ServiceInstance serviceInstance : serviceInstances) {
                String instanceZone = getZone(serviceInstance);
                if (zone.equalsIgnoreCase(instanceZone)) {
                    filteredInstances.add(serviceInstance);
                }
            }
            if (filteredInstances.size() > 0) {
                return filteredInstances;
            }
        }
        return serviceInstances;
    }

    private String getZone(ServiceInstance serviceInstance) {
        Map<String, String> metadata = serviceInstance.getMetadata();
        if (metadata != null) {
            return metadata.get(ZONE);
        }
        return null;
    }

    @Override
    public Flux<List<ServiceInstance>> get(Request request) {
        String zone = getZone(request);
        return get().map(serviceInstances -> filteredByZone(zone, serviceInstances));
    }

    private String getZone(Request request) {
        String zone = null;
        //通过ReactiveLoadBalancerClientFilter，调用时可以从这里获取request
        if (request instanceof DefaultRequest) {
            DefaultRequest defaultRequest = (DefaultRequest) request;
            if (defaultRequest.getContext() instanceof RequestDataContext) {
                RequestDataContext requestDataContext = (RequestDataContext) defaultRequest.getContext();
                List<String> zones = requestDataContext.getClientRequest().getHeaders().get(ZONE);
                if (zones != null && zones.size() > 0) {
                    zone = zones.get(0);
                }
            }
        }
        if (StringUtils.isNotEmpty(zone)) {
            return zone;
        }
        //通过ThreadLocal获取，异步时可能获取不到
        HttpServletRequest originServletRequest = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        if (originServletRequest != null) {
            zone = originServletRequest.getHeader(ZONE);
        }
        return zone;
    }


}
